package DiffieHellman;

import NumTheory.NumTheory;

import java.math.BigInteger;
import java.util.Random;

public class DHexchange {

    private static final Random RND = new Random();

    private static final NumTheory NUM_THEORY = new NumTheory();

    // 选择的大素数
    private BigInteger bigPrime;
    // GF(p)的本原根
    private BigInteger originRoot;

    public BigInteger pka;
    public BigInteger pkb;
    public BigInteger keya;
    public BigInteger keyb;

    public DHexchange(int parameter) {
        BigInteger one = BigInteger.ONE;
        BigInteger two = BigInteger.valueOf(2);
        // g是m的原根的充分必要条件是：φ(m)的所有的不同的素因子qi都满足 g^(φ(m)/qi) ≠ 1 mod m。
        // 因为p=2q+1，φ(p)=p-1=2q，不同的素因子只有2和q，所以均不同余于1，可以说明g是p的原根
        BigInteger temp = null;
        do {
            temp = BigInteger.probablePrime(parameter, RND);
            bigPrime = temp.multiply(two).add(one);
        } while (!NUM_THEORY.MRtest(bigPrime, 5));

        // 这里 bigPrime 是素数，且φ(bigPrime)的素因子只有 2 和 temp
        do {
            originRoot = BigInteger.probablePrime(RND.nextInt(parameter), RND);
        } while (originRoot.modPow(two, bigPrime).equals(one) ||
                originRoot.modPow(temp, bigPrime).equals(one));
    }

    public DHexchange(BigInteger bigPrime, BigInteger originRoot) {
        this.bigPrime = bigPrime;
        this.originRoot = originRoot;
    }

    public boolean DH(BigInteger skA, BigInteger skB) {

        // 用户A生成私钥skA，计算pka
        pka = NUM_THEORY.fastPow(originRoot, skA, bigPrime);
        // 用户B生成私钥skB，计算pkb
        pkb = NUM_THEORY.fastPow(originRoot, skB, bigPrime);

        // 互相交换pkA和pkB，分别计算私钥
        keya = NUM_THEORY.fastPow(pkb, skA, bigPrime);
        keyb = NUM_THEORY.fastPow(pka, skB, bigPrime);

        // A和B分别在本地计算出了keyFromA和keyFromB，如果相同则成功
        return keya.equals(keyb);
    }
}
